﻿extern "C" {
#include <libavutil/opt.h>
#include <libavutil/samplefmt.h>
#include <libavutil/channel_layout.h>
#include <libavutil/common.h>
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavfilter/avfilter.h>
#include <libavfilter/buffersink.h>
#include <libavfilter/buffersrc.h>
#include <libavutil/bprint.h>
}
#include <iostream>

#define INPUT_SAMPLERATE 32000
#define INPUT_FORMAT AV_SAMPLE_FMT_S16
#define INPUT_CHANNEL_LAYOUT AV_CH_LAYOUT_MONO

int main()
{
    const char* input_filename = "input.wav";
    const char* output_filename = "output.jpg";

    av_register_all();
    avcodec_register_all();
    avfilter_register_all();

    AVFormatContext* format_ctx = avformat_alloc_context();
    if (avformat_open_input(&format_ctx, input_filename, NULL, NULL) != 0)
    {
        fprintf(stderr, "Error opening input file\n");
        return -1;
    }

    if (avformat_find_stream_info(format_ctx, NULL) < 0)
    {
        fprintf(stderr, "Could not find stream information\n");
        return -1;
    }

    AVCodec* codec = avcodec_find_decoder(format_ctx->streams[0]->codecpar->codec_id);
    AVCodecContext* codec_ctx = avcodec_alloc_context3(codec);
    avcodec_parameters_to_context(codec_ctx, format_ctx->streams[0]->codecpar);
    avcodec_open2(codec_ctx, codec, NULL);

    AVFilterGraph* filter_graph = avfilter_graph_alloc();
    AVFilterContext* buffersrc_ctx = NULL;
    AVFilterContext* showspectrum_ctx = NULL;
    AVFilterContext* buffersink_ctx = NULL;

    const AVFilter* abuffer = avfilter_get_by_name("abuffer");
    const AVFilter* showspectrum = avfilter_get_by_name("showspectrum");
    const AVFilter* buffersink = avfilter_get_by_name("buffersink");

    AVBPrint args;
    av_bprint_init(&args, 0, AV_BPRINT_SIZE_AUTOMATIC);
    av_bprintf(&args, "time_base=%d/%d:sample_rate=%d:sample_fmt=%s",
        1, codec_ctx->sample_rate,
        codec_ctx->sample_rate,
        av_get_sample_fmt_name(codec_ctx->sample_fmt));

    char ch_layout[64] = { 0 };
    av_get_channel_layout_string(ch_layout, sizeof(ch_layout), 0, AV_CH_LAYOUT_MONO);
    av_bprintf(&args, ":channel_layout=%s",
        ch_layout);

    avfilter_graph_create_filter(&buffersrc_ctx, abuffer, "in", args.str, NULL, filter_graph);
    avfilter_graph_create_filter(&buffersink_ctx, buffersink, "out", NULL, NULL, filter_graph);

    const char* show_filter_desc = "s=1480x680:scale=5thrt:color=intensity:legend=0:start=20:stop=16000";
    //const char* show_filter_desc = "s=480x480:scale=5thrt:saturation=10:color=fire:legend=0:start=20:stop=16000";
    avfilter_graph_create_filter(&showspectrum_ctx, showspectrum, "showspectrum", show_filter_desc, NULL, filter_graph);


    avfilter_link(buffersrc_ctx, 0, showspectrum_ctx, 0);
    avfilter_link(showspectrum_ctx, 0, buffersink_ctx, 0);

    avfilter_graph_config(filter_graph, NULL);

    // 打开输出文件
    //FILE* output_file = fopen(output_filename, "wb");

    AVFrame* frame = av_frame_alloc();
    AVPacket packet;
    av_init_packet(&packet);

    int i = 0;
    while (av_read_frame(format_ctx, &packet) >= 0)
    {
        if (packet.stream_index == 0)
        {
            if (avcodec_send_packet(codec_ctx, &packet) == 0)
            {
                while (avcodec_receive_frame(codec_ctx, frame) == 0)
                {
                    // 发送音频帧到滤镜图
                    av_buffersrc_add_frame_flags(buffersrc_ctx, frame, AV_BUFFERSRC_FLAG_KEEP_REF);

                    // 获取输出帧
                    AVFrame* filtered_frame = av_frame_alloc();
                    int ret = av_buffersink_get_frame(buffersink_ctx, filtered_frame);
                    if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
                    {
                        av_frame_free(&filtered_frame);
                        continue;
                    }
                    else if (ret < 0)
                    {
                        fprintf(stderr, "Error getting frame from filter\n");
                        break;
                    }

                    // 将输出帧转换为图像
                    AVCodec* output_codec = avcodec_find_encoder(AV_CODEC_ID_PNG);
                    AVCodecContext* output_codec_ctx = avcodec_alloc_context3(output_codec);
                    output_codec_ctx->width = filtered_frame->width;
                    output_codec_ctx->height = filtered_frame->height;
                    output_codec_ctx->pix_fmt = AV_PIX_FMT_RGBA;
                    // Set the timebase
                    output_codec_ctx->time_base ={ 1, codec_ctx->time_base.den };; // for example, 1/25

                    //std::cout << "output_codec_ctx->output_codec_ctx.den:" << output_codec_ctx->time_base.den;
                    //std::cout << ",output_codec_ctx->output_codec_ctx.num:" << output_codec_ctx->time_base.num << std::endl;

                    avcodec_open2(output_codec_ctx, output_codec, NULL);

                    
                    AVPacket output_packet;
                    av_init_packet(&output_packet);
                    avcodec_send_frame(output_codec_ctx, filtered_frame);
                    avcodec_receive_packet(output_codec_ctx, &output_packet);

                    // 写入输出文件
                    FILE* output_file = fopen(output_filename, "wb+");
                    fwrite(output_packet.data, 1, output_packet.size, output_file);
                    fclose(output_file);


                    av_packet_unref(&output_packet);
                    avcodec_free_context(&output_codec_ctx);
                    av_frame_free(&filtered_frame);

                    // 将输出帧写入文件
                    //fwrite(filtered_frame->data[0], 1, filtered_frame->linesize[0], output_file);

                    av_frame_free(&filtered_frame);
                }
            }
        }
        av_packet_unref(&packet);
    }

    //fclose(output_file);

    avformat_close_input(&format_ctx);
    avcodec_free_context(&codec_ctx);
    avfilter_free(buffersrc_ctx);
    avfilter_free(buffersink_ctx);
    avfilter_graph_free(&filter_graph);
    av_frame_free(&frame);

    avformat_network_deinit();

    return 0;
}
